
import numpy as np

import os
import pickle

import torch

import shutil



def save_mesh(out_file, vertices):
    if type(vertices) == torch.Tensor:
        vertices = vertices.squeeze().cpu().numpy()
    model_file = '/home/ljl/lidarcap/data/basicModel_neutral_lbs_10_207_0_v1.0.0.pkl'
    with open(model_file, 'rb') as f:
        smpl_model = pickle.load(f, encoding='iso-8859-1')
        face_index = smpl_model['f'].astype(np.int64)
    face_1 = np.ones((face_index.shape[0], 1))
    face_1 *= 3
    face = np.hstack((face_1, face_index)).astype(int)
    if os.path.exists(out_file):
        os.remove(out_file)
    with open(out_file, "wb") as zjy_f:
        np.savetxt(zjy_f, vertices, fmt='%f %f %f')
        np.savetxt(zjy_f, face, fmt='%d %d %d %d')
    ply_header = '''ply
format ascii 1.0
element vertex 6890
property float x
property float y
property float z
element face 13776
property list uchar int vertex_indices
end_header
    '''
    with open(out_file, 'r+') as f:
        old = f.read()
        f.seek(0)
        f.write(ply_header)
        f.write(old)

if __name__ == '__main__':

    smpl_path = [0 for i in range(6)]
    smpl_path[0] = '/home/ljl/hdd/lidarcap/labels/3d/smpl/39/'
    smpl_path[1] = '/home/ljl/hdd/lidarcap/labels/3d/smpl/33/'
    smpl_path[2] = '/home/ljl/hdd/lidarcap/labels/3d/smpl/35/'
    smpl_path[3] = '/home/ljl/hdd/lidarcap/labels/3d/smpl/27/'
    smpl_path[4] = '/home/ljl/hdd/lidarcap/labels/3d/smpl/32/'
    smpl_path[5] = '/home/ljl/hdd/lidarcap/labels/3d/smpl/31/'
    save_path = '/home/ljl/hdd/teaser_dataset/'
    filenames = [0 for i in range(6)]
    filenames[0] = [2113, 3414, 5127, 6001, 9234, 11938]
    filenames[1] = [2076, 2301, 2684, 4080, 4477, 5049]
    filenames[2] = [6380, 7510, 8383, 8803, 8994, 16503]
    filenames[3] = [340, 556, 728, 919, 1921, 2401]
    filenames[4] = [497, 599, 881, 1379, 4392, 5896]
    filenames[5] = [2987, 3583, 4998, 6405, 6764, 8880]

    if os.path.exists(save_path):
        shutil.rmtree(save_path)

    for i in range(6):
        out_path = save_path + str(i) + "/"
        os.makedirs(out_path)
        for j in range(len(filenames[i])):
            filename = filenames[i][j]
            file_path = smpl_path[i] + str(filename) +'.npz'
            npz_file = np.load(file_path)
            vertices = npz_file['vertices']
            vertices += np.array([3 * j, 1.5 * i, 0])
            out_file = out_path + str(filename) + '.ply'
            save_mesh(out_file, vertices)








